using System;
using System.Reflection;
using System.Collections;
using NBody;
using NBody.Analysis;
using NBody.Cosmology;
using OptionParser;
using DataFile;

[assembly: AssemblyTitle ("NBodyVelDisp")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Compute various velocity dispersion quantities.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class VelDispOptions : InputOutputGetOptions
{
    [Option ("-b", "Number of bins")]
    public int NumBins;
    
    [Option ("-s", "Bin spacing")]
    public Bin.Spacing Spacing;
    
    [Option ("-t", "Output total dispersions")]
    public bool TotalDisp;
    
    [Option ("-c", "Adjust to centre of mass")]
    public bool CentreOfMass;
    
    [Option ("-r", "Output in relative units")]
    public bool Relative;
    
    public VelDispOptions()
    {
        NumBins = 100;
        Spacing = Bin.Spacing.Logarithmic;
    }
}

public class NBodyVelDisp
{
    static double r_200 = 1.0;
    static double sigma_kms = 1.0;
    
    public static void Main(string[] args)
    {
        VelDispOptions opts = new VelDispOptions();
        opts.ProcessArgs(args, ".veldisp");
                        
        NBodySystem s = NBodySystem.Read(opts.InFile);
        
        if (opts.CentreOfMass)
            s.AdjustForCentreOfMass(1e-3);
            
        if (opts.Relative)
        {
            Model cosmo = Model.SCDM;
            double z = cosmo.Redshift(s.Time);
            double rho_c = cosmo.Density(z);
            r_200 = VirialRadius.CalculateR200(s, rho_c);
            sigma_kms = 20.74 * 20.74;
        }
        
        if (opts.TotalDisp)
        {
            Console.Error.WriteLine("Velocity dispersions:");
            double s2_r = CalculateTotalDispersion(s, new Bin.Function(vr));
            Console.WriteLine("Radial: {0}", s2_r);
            double s2_t = CalculateTotalDispersion(s, new Bin.Function(vt));
            Console.WriteLine("Theta: {0}", s2_t);
            double s2_p = CalculateTotalDispersion(s, new Bin.Function(vp));
            Console.WriteLine("Phi: {0}", s2_p);
            
            Console.WriteLine("Total: {0}", 1.0/3.0 * (s2_r + s2_t + s2_p));
            Console.WriteLine("Beta: {0}", 1.0 - (s2_t + s2_p) / 2.0 / s2_r);
            Environment.Exit(0);
        }
        
        Table data = new Table(opts.NumBins, 4, "Radius, sigma^2_r, sigma^2_tan, beta");
        data.Clear();
        
        VelocityBin sigma2_r = new VelocityBinDispersion();
        sigma2_r.BinSpacing = opts.Spacing;
        sigma2_r.NumBins = opts.NumBins;
        sigma2_r.VelocityDirection = VelocityBin.Direction.Radial;
        sigma2_r.Analyse(s);
        
        VelocityBin sigma2_t = new VelocityBinDispersion();
        sigma2_t.BinSpacing = opts.Spacing;
        sigma2_t.NumBins = opts.NumBins;
        sigma2_t.VelocityDirection = VelocityBin.Direction.Theta;
        sigma2_t.Analyse(s);
        
        VelocityBin sigma2_p = new VelocityBinDispersion();
        sigma2_p.BinSpacing = opts.Spacing;
        sigma2_p.NumBins = opts.NumBins;
        sigma2_p.VelocityDirection = VelocityBin.Direction.Phi;
        sigma2_p.Analyse(s);
        
        for (int i = 0; i < data.Length; i++)
        {
            data[i, 0] = sigma2_r.Data[i, 0];
            data[i, 1] = Math.Sqrt(sigma2_r.Data[i, 1]);
            data[i, 2] = Math.Sqrt(1.0/2.0 * (sigma2_t.Data[i, 1] + sigma2_p.Data[i, 1]));
            data[i, 3] = 1.0 - (sigma2_t.Data[i, 1] + sigma2_p.Data[i, 1]) / 2.0 / sigma2_r.Data[i, 1];
        }
        
        data.Print(opts.OutFile, new DataFile.Transform[] { RelRad, RelSigmaR, RelSigmaT, DataFile.Table.IdentityTransform });
    }
    
    static double vr(Particle p)
    {
        return p.RadialVelocity;
    }
    
    static double vt(Particle p)
    {
        return p.ThetaVelocity;
    }
    
    static double vp(Particle p)
    {
        return p.PhiVelocity;
    }
    
    static double RelRad(double r)
    {
        return Math.Log10(r / r_200);
    }
    
    static double RelSigmaR(double s)
    {
        return s * sigma_kms;
    }
    
    static double RelSigmaT(double s)
    {
        return s * sigma_kms;
    }
    
    static double CalculateTotalDispersion(NBodySystem s, Bin.Function f)
    {
        double v = 0.0, v2 = 0.0;
        foreach (Particle p in s)
        {
            double t = f(p);
            v += t;
            v2 += t * t;
        }
        double n = (double)s.NumParts;
        double vbar = v / n;
        double v2bar = v2 / n;
        
        return v2bar - (vbar * vbar);
    }
}
